// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/memory/memory.h>
#include <pollux/exec/operator.h>
#include <pollux/plan/plan_builder.h>
#include <pollux/vector/base_vector.h>

// Test utilities, for convenience and conciseness.
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/vector/vector_maker.h>

#include <melon/init/init.h>
#include <algorithm>

using namespace kumo::pollux;
using namespace kumo::pollux::plan;

// This file contains a step-by-step example of a custom user-defined operator
// using the operator extensibility API. It defines a new trivial
// "DuplicateRowOperator" that simply duplicates every input row in the output.
// It does so by wrapping inputs in a dictionary vector to increase its
// cardinality in a zero-copy manner.

// In order to define a new operator, four steps required:
//
// #1. Define a new custom plan node.
// #2. Define the new operator, extending the base exec::Operator class.
// #3. Define the plan translation logic (create operator based on plan node).
// #4. Register the plan translator.

// First, let's define a custom plan node:
class DuplicateRowNode : public core::PlanNode {
public:
    // Our example takes a single source.
    DuplicateRowNode(const core::PlanNodeId &id, core::PlanNodePtr source)
        : PlanNode(id), sources_{std::move(source)} {
    }

    // Output is always the exact same as the single input source.
    const RowTypePtr &outputType() const override {
        return sources_.front()->outputType();
    }

    const std::vector<core::PlanNodePtr> &sources() const override {
        return sources_;
    }

    std::string_view name() const override {
        return "duplicate row";
    }

private:
    // One can add details about the plan node and its metadata in a textual
    // format.
    void addDetails(std::stringstream & /* stream */) const override {
    }

    std::vector<core::PlanNodePtr> sources_;
};

// Second, let's define the operator. Here's where the main logic lives.
class DuplicateRowOperator : public exec::Operator {
public:
    // The operator takes the plan node defined above plus additional env
    // parameters for execution.
    DuplicateRowOperator(
        int32_t operatorId,
        exec::DriverCtx *driverCtx,
        const std::shared_ptr<const DuplicateRowNode> &planNode)
        : Operator(
            driverCtx,
            nullptr,
            operatorId,
            planNode->id(),
            "DuplicateRow") {
    }

    // Called every time there's input available. We just save it in the `input_`
    // member defined in the base class, and process it on `getOutput()`.
    void addInput(RowVectorPtr input) override {
        input_ = input;
    }

    bool needsInput() const override {
        return !noMoreInput_;
    }

    // Called every time your operator needs to produce data. It processes the
    // input saved in `input_`, wraps it in a dictionary, and returns a new
    // RowVector.
    RowVectorPtr getOutput() override {
        if (!input_) {
            return nullptr;
        }

        // We move `input_` to signal the input has been processed.
        auto input = std::move(input_);

        // Creates dictionary indices with twice the size, and set duplicated
        // consecutive indices (0, 0, 1, 1, 2, 2, ...). This allows us to increase
        // the cardinality and produce output without touching the actual data. The
        // assumption is that the data could be complex/nested data types which
        // could be expensive to copy.
        const size_t outputSize = input->size() * 2;
        BufferPtr indices = AlignedBuffer::allocate<int32_t>(outputSize, pool());
        auto rawIndices = indices->asMutable<int32_t>();

        for (size_t i = 0; i < outputSize; ++i) {
            rawIndices[i] = i / 2;
        }

        // Wrap each column in a dictionary, create a new row vector, and return.
        std::vector<VectorPtr> outputChildren;
        outputChildren.reserve(input->childrenSize());

        for (const auto &child: input->children()) {
            outputChildren.push_back(BaseVector::wrap_in_dictionary(
                BufferPtr(), indices, outputSize, child));
        }
        return std::make_shared<RowVector>(
            pool(),
            input->type(),
            BufferPtr(),
            outputSize,
            std::move(outputChildren));
    }

    // This simple operator is never blocked.
    exec::BlockingReason isBlocked(ContinueFuture *future) override {
        return exec::BlockingReason::kNotBlocked;
    }

    bool isFinished() override {
        return noMoreInput_;
    }
};

// Third, we need to define a plan translation logic to convert our custom plan
// node into our custom operator. Check `pollux/exec/LocalPlanner.cpp` for more
// details.
class DuplicateRowTranslator : public exec::Operator::PlanNodeTranslator {
    std::unique_ptr<exec::Operator>
    toOperator(exec::DriverCtx *ctx, int32_t id, const core::PlanNodePtr &node) {
        if (auto dupRowNode =
                std::dynamic_pointer_cast<const DuplicateRowNode>(node)) {
            return std::make_unique<DuplicateRowOperator>(id, ctx, dupRowNode);
        }
        return nullptr;
    }
};

int main(int argc, char **argv) {
    melon::Init init{&argc, &argv};

    // Fourth, we register the custom plan translator. We're now ready to use our
    // operator in a query plan.
    exec::Operator::registerOperator(std::make_unique<DuplicateRowTranslator>());

    // From now on we will create a query plan and input dataset, execute it, an
    // assert that the output results contain the dataset properly duplicated.

    // Create a new memory pool to in this example.
    memory::MemoryManager::initialize({});
    auto pool = memory::memoryManager()->addLeafPool();

    // VectorMaker is a test utility that helps you build vectors. Shouldn't be
    // used in production.
    VectorMaker maker{pool.get()};

    // Create an input dataset containing two unnamed columns (INTEGER and
    // VARCHAR), and 5 records:
    auto inputRowVector = maker.row_vector({
        maker.flat_vector({0, 1, 2, 3, 4}),
        maker.flat_vector({"a", "b", "c", "d", "e"}),
    });

    // Create a query plan containing a ValuesNode (to let you pump input datasets
    // directly into the operator chain), and our custom plan node.
    auto plan = PlanBuilder()
            .values({inputRowVector})
            .addNode([](std::string id, core::PlanNodePtr input) {
                return std::make_shared<DuplicateRowNode>(id, input);
            })
            .planNode();

    // Create the expected results to ensure our plan produces the correct
    // results.
    auto expectedResults = maker.row_vector({
        maker.flat_vector({0, 0, 1, 1, 2, 2, 3, 3, 4, 4}),
        maker.flat_vector({"a", "a", "b", "b", "c", "c", "d", "d", "e", "e"}),
    });

    // Execute the plan above and assert it produces the results in
    // `expectedResults`.
    exec::test::AssertQueryBuilder(plan).assertResults(expectedResults);

    return 0;
}
